package main

import (
	"fmt"
	"go/ast"
	"log"
	"os"
	"strings"
)

type traceParam struct {
	typ    string
	params []string
}

func (g *generator) traceGetType(raw ast.Expr, imports map[string]string) (typ string) {
	switch t := raw.(type) {
	case *ast.Ident:
		typ = t.Name
	case *ast.StarExpr:
		switch tt := t.X.(type) {
		case *ast.SelectorExpr:
			// 其他包定义的指针类型
			typ = fmt.Sprintf("*%s.%s", tt.X, tt.Sel)
			// 导入其他包的定义
			g.pickupImports(fmt.Sprint(tt.X), imports)
		case *ast.Ident:
			// 当前包定义的指针类型
			typ = fmt.Sprintf("*%s", tt.Name)
		default:
			log.Printf("not supported type: %#v\n", tt)
		}
	case *ast.SelectorExpr:
		typ = fmt.Sprintf("%s.%s", t.X, t.Sel)
		g.pickupImports(fmt.Sprint(t.X), imports)
	case *ast.ArrayType:
		atyp := g.traceGetType(t.Elt, imports)
		if atyp != "" {
			typ = fmt.Sprintf("[]%s", atyp)
		}
	default:
	}
	return
}

func (g *generator) traceGetFuncDecl(fields []*ast.Field, imports map[string]string) (decl string, call string) {
	var traceParams []*traceParam
	for _, parameter := range fields {
		typ := g.traceGetType(parameter.Type, imports)
		if typ == "" {
			log.Println("unsupported parameter type.")
			return
		}
		tp := &traceParam{
			typ: typ,
		}
		for _, name := range parameter.Names {
			tp.params = append(tp.params, name.Name)
		}
		traceParams = append(traceParams, tp)
	}
	var plist []string
	var nlist []string
	for _, tp := range traceParams {
		nlist = append(nlist, tp.params...)
		plist = append(plist, strings.Join(tp.params, ",")+" "+tp.typ)
	}
	return strings.Join(plist, ","), strings.Join(nlist, ",")
}

func (g *generator) generateTrace(gt *genTag, imports map[string]string) {
	fn, ok := gt.Decl.(*ast.FuncDecl)
	if !ok {
		return
	}
	if isExported(fn.Name.Name) {
		return
	}
	if fn.Recv != nil {
		log.Println("func type not supported.")
		return
	}
	// 检查第一个参数是否ctx
	if len(fn.Type.Params.List) == 0 {
		return
	}
	if len(fn.Type.Params.List[0].Names) != 1 {
		return
	}
	if fn.Type.Params.List[0].Names[0].Name != "ctx" {
		return
	}

	pd, pc := g.traceGetFuncDecl(fn.Type.Params.List[1:], imports)
	rd, _ := g.traceGetFuncDecl(fn.Type.Results.List, imports)
	if strings.Contains(rd, ",") {
		rd = "(" + rd + ")"
	}

	pkgName, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	index := strings.Index(pkgName, "src/ndp/")
	if index > 0 {
		pkgName = pkgName[index+4:]
	}
	if pd == "" {
		g.p("func ", strings.Title(fn.Name.Name), "(ctx context.Context) ", rd, "{")
	} else {
		g.p("func ", strings.Title(fn.Name.Name), "(ctx context.Context, ", pd, ") ", rd, "{")
	}
	g.p(`ctx = trace.WithLables(ctx, "`, pkgName+"."+strings.Title(fn.Name.Name), `")`)
	if len(rd) > 0 {
		g.p("return ", fn.Name.Name, "(ctx, ", pc, ")")
	} else {
		g.p(fn.Name.Name, "(ctx, ", pc, ")")
	}
	g.p("}")
	g.p()
}
